import asyncio
import logging
import random

import settings
from common.information import port_dict
from common.sys_settings import sysSetting
from common.data import net
from settings import server
import time
import functools


async def pipe(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, _client=False):
    """
    管道,实现内容的转发
    """
    try:
        while not reader.at_eof():
            data = await reader.read(sysSetting.bufsize)
            writer.write(data)
            await writer.drain()
    finally:
        logging.info('管道断开')
        if not _client:
            writer.close()


async def create_pipe(tunnel_reader, tunnel_writer, uid, vid):
    if vid not in visitors[uid]:
        logging.warning(f"用户-{vid} 不存在")
        return
    v_reader, v_writer = visitors[uid][vid]

    task_list = [
        asyncio.create_task(pipe(v_reader, tunnel_writer, True)),
        asyncio.create_task(pipe(tunnel_reader, v_writer))
    ]

    await asyncio.wait(task_list)
    try:
        del visitors[uid][vid]
    except KeyError as e:
        logging.error(f"没有该用户{vid}或客户端{uid},异常为{e}")
    _message = net.encode(f'用户{vid}断开连接', sysSetting.message, 'utf-8', sysSetting.separator)
    tunnel_writer.write(_message)
    await tunnel_writer.drain()
    logging.info(f"用户{vid}断开连接")
    tunnel_writer.close()


async def tcpserver_handle(tunnel_writer: asyncio.StreamWriter, uid,
                           reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    while True:
        vid = random.randint(1, 10086)
        if vid not in visitors[uid]:
            break
    visitors[uid][vid] = (reader, writer)
    _message = net.encode(_data=[vid, writer.get_extra_info('peername')],
                          _type_of_message=sysSetting.new_pipe,
                          encoding='utf-8',
                          separator=sysSetting.separator
                          )

    # 向客户端发送构建管道的请求
    logging.info(f"{visitors[uid][vid][1].get_extra_info('peername')}({vid})接入网络")
    try:
        tunnel_writer.write(_message)
        await tunnel_writer.drain()
    except Exception as e:
        logging.error(f'{visitors[uid][vid][1].get_extra_info("peername")}({vid})在连接服务器-{uid}时发生了一个错误错误:{e}')
        try:
            del visitors[uid][vid]
        except KeyError as ke:
            logging.error(f"{ke}")


async def heart_beat_is_overtime(writer, last_heartbeat):
    """
    连接超时后会抛出异常
    """
    while True:
        if time.time() - last_heartbeat[0] > sysSetting.max_heart_beat_time:
            writer.close()
            raise asyncio.exceptions.TimeoutError
        await asyncio.sleep(10)


async def keep_alive_control(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, uid: int):
    """
    保证连接
    """
    last_heartbeat = [time.time(), 0]
    _loop = 0
    heartbeat_task = asyncio.create_task(heart_beat_is_overtime(writer, last_heartbeat))
    try:
        while True:
            _message = await reader.readuntil(sysSetting.separator)
            _type, _data = net.decode(_message, 'utf-8', sysSetting.separator)
            if _type == sysSetting.ping:
                _type = sysSetting.pong
                _data = ''
                _message = net.encode(_data, _type, 'utf-8', sysSetting.separator)
                if _loop % sysSetting.when_log_ping == 0:
                    logging.info(f'来自客户端-{uid} ping({_loop})')
                    if _loop >= sysSetting.max_ping_loop:
                        _loop = 0
                _loop = _loop + 1
                writer.write(_message)
                await writer.drain()
                last_heartbeat[0] = time.time()
    except asyncio.IncompleteReadError as e:
        logging.error(f'客户端-{uid}连接发生异常:{e}')
    finally:
        heartbeat_task.cancel()
        writer.close()
        try:
            del client[uid]
            del visitors[uid]
        except KeyError as e:
            logging.error(f"{e}")
        logging.info(f'客户端-{uid}连接已关闭')


async def start_tcp_server(reader: asyncio.StreamReader, writer: asyncio.StreamWriter, uid):
    """
        开启转发服务器
    """

    port = port_dict_obj.get_port()
    if port == 0:
        _type = sysSetting.error
        _data = '无可用端口'
        _message = net.encode(_data, _type, 'utf-8', sysSetting.separator)
        writer.write(_message)
        await writer.drain()
        writer.close()
        try:
            del client[uid]
            del visitors[uid]
        except KeyError as e:
            logging.error(f"{e}")
        logging.info(f"无可用端口,服务器-{uid}关闭")
        return
    port_dict_obj.change(port, True)
    try:

        tcp_server = await asyncio.start_server(
            functools.partial(tcpserver_handle, writer, uid),
            # 这一段代码很关键否则会导致连接出问题
            # 为什么这段代码得到的writer对象不同,在回调后,writer对象被设置成了future pending,而在传参前并没有
            # waiter=<Future pending cb=[Task.task_wakeup()]>
            host=server['remote']['ip'],
            port=port
        )
        await tcp_server.start_serving()

    except Exception as e:
        port_dict_obj.change(port, False)
        logging.error(f'发生错误:{e}')
        visitors.clear()
        logging.info(f'服务器-{uid}发生错误并关闭')
        _message = net.encode(_data=e, _type_of_message=sysSetting.error, encoding='utf-8',
                              separator=sysSetting.separator)
        writer.write(_message)
        await writer.drain()
        try:
            del client[uid]
            del visitors[uid]
        except KeyError as e:
            logging.error(f"{e}")
        writer.close()
        return

    _type = sysSetting.get_host
    _data = (port, uid)
    _message = net.encode(_data, _type, 'utf-8', sysSetting.separator)
    writer.write(_message)
    await writer.drain()
    logging.info(f'类型{_type} 服务器-{uid}绑定了{port}开启')

    task = [
        asyncio.create_task(keep_alive_control(reader, writer, uid))
    ]

    try:

        await asyncio.wait(task, return_when=asyncio.FIRST_EXCEPTION)
    finally:
        for t in task:
            t.cancel()
        port_dict_obj.change(port, False)
        try:
            del client[uid]
            del visitors[uid]
        except KeyError as e:
            logging.error(f"{e}")
        writer.close()
        tcp_server.close()
        logging.info(f'客户端-{uid}断开连接')


async def send_visitor_message(writer: asyncio.StreamWriter, uid):
    connectors = {}
    for connector in visitors[uid]:
        connectors[connector] = visitors[uid][connector][1].get_extra_info('peername')
    _message = net.encode(_data=f'用户列表:{connectors}', _type_of_message=sysSetting.message, encoding='utf-8',
                          separator=sysSetting.separator)
    writer.write(_message)
    await writer.drain()
    writer.close()


async def main_server_handle(reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
    """
        建立隧道 ,接收客户端的指令
    """
    message = await reader.readuntil(sysSetting.separator)
    _type, _data = net.decode(_message=message, separator=sysSetting.separator, encoding='utf-8')
    logging.info(f"类型:{_type}")
    if _type == sysSetting.new_server:
        # 开启一个新的服务器
        while True:
            uid = random.randint(1, 10086)
            if uid not in client:
                break
        client[uid] = writer.get_extra_info('peername')
        visitors[uid] = {}
        return await start_tcp_server(reader, writer, uid)
    elif _type == sysSetting.new_pipe:
        # 开启新的管道,用来连接远端用户和内网服务器
        return await create_pipe(reader, writer, _data[0], _data[1])
    elif _type == sysSetting.request_get_connector:
        return await send_visitor_message(writer, _data)


async def run():
    tcp_server = await asyncio.start_server(main_server_handle,
                                            host=server['remote']['ip'],
                                            port=server['remote']['port']
                                            )
    await tcp_server.serve_forever()


def main():
    asyncio.run(run())


if __name__ == '__main__':
    fmt = '%(asctime)s - %(levelname)s: %(message)s'
    logging.basicConfig(format=fmt, level=logging.DEBUG)
    client = {}  # uid:peername
    visitors = {}  # {uid:{vid:reader,writer}}
    port_dict_obj = port_dict(settings.max_connect, settings.d_port, settings.open_port)
    port_dict_obj.init_port_dict()
    logging.info('端口加载完成')
    logging.info('开启服务')
    main()
